# -*- coding: binary -*-

#
# This class implements an override for RubySMB's default authentication method to instead
# use a kerberos authenticator
#
module Msf::Exploit::Remote::DCERPC::KerberosAuthentication
  # @param [Msf::Exploit::Remote::Kerberos::ServiceAuthenticator::SMB] kerberos_authenticator The authenticator to make the required Kerberos requests
  def kerberos_authenticator=(kerberos_authenticator)
    @kerberos_authenticator = kerberos_authenticator
  end

  # Initialize the auth provider using Kerberos
  # @return Serialized message for initializing the auth provider
  def auth_provider_init
    kerberos_result = @kerberos_authenticator.authenticate
    @application_key = @session_key = kerberos_result[:session_key]
    @client_sequence_number = kerberos_result[:client_sequence_number]
    kerberos_result[:security_blob]
  end

  # Encrypt the value in dcerpc_req, and add a valid signature to the request.
  # This function modifies the request object in-place, and does not return anything.
  # @param dcerpc_req [Request] The Request object to be encrypted and signed in-place
  def auth_provider_encrypt_and_sign(dcerpc_req)
    auth_pad_length = get_auth_padding_length(dcerpc_req.stub.to_binary_s.length)
    plain_stub = dcerpc_req.stub.to_binary_s + "\x00" * auth_pad_length
    emessage, header_length, krb_pad_length = self.krb_encryptor.encrypt_and_increment(plain_stub)

    encrypted_stub = emessage[header_length..-1]
    signature = emessage[0,header_length]
    set_encrypted_packet(dcerpc_req, encrypted_stub, auth_pad_length)
    set_signature_on_packet(dcerpc_req, signature)
  end

  # Decrypt the value in dcerpc_response, and validate its signature.
  # This function modifies the request object in-place, and returns whether the signature was valid.
  # @param dcerpc_response [Response] The Response packet to decrypt and verify in-place
  # @raise ArgumentError If the auth type is not SPNEGO (which ultimately wraps Kerberos)
  # @return [Boolean] Is the packet's signature valid?
  def auth_provider_decrypt_and_verify(dcerpc_response)
    auth_type = dcerpc_response.sec_trailer.auth_type
    unless [RubySMB::Dcerpc::RPC_C_AUTHN_GSS_NEGOTIATE].include?(auth_type)
      raise ArgumentError, "Unsupported Auth Type: #{dcerpc_response.sec_trailer.auth_type}"
    end
    encrypted_stub = get_response_full_stub(dcerpc_response)
    signature = dcerpc_response.auth_value
    data = signature + encrypted_stub

    begin
      result = self.krb_encryptor.decrypt_and_verify(data)
    rescue Rex::Proto::Kerberos::Model::Error::KerberosError
      return false
    end
    set_decrypted_packet(dcerpc_response, result)

    true
  end

  def build_ap_rep(session_key, sequence_number)
    pvno = Rex::Proto::Kerberos::Model::VERSION
    msg_type = Rex::Proto::Kerberos::Model::AP_REP
    ctime = Time.now.utc
    cusec = ctime&.usec

    encrypted_part = Rex::Proto::Kerberos::Model::EncApRepPart.new(
      ctime: ctime,
      cusec: cusec,
      sequence_number: sequence_number,
      enc_key_usage: Rex::Proto::Kerberos::Crypto::KeyUsage::AP_REP_ENCPART
    )
    enc_aprep = Rex::Proto::Kerberos::Model::EncryptedData.new(
      etype: session_key.type,
      cipher: encrypted_part.encrypt(session_key.type, session_key.value)
    )

    Rex::Proto::Kerberos::Model::ApRep.new(
      pvno: pvno,
      msg_type: msg_type,
      enc_part: enc_aprep
    )
  end

  def auth_provider_complete_handshake(response, options)
    begin
      @kerberos_authenticator.validate_response!(response.auth_value, accept_incomplete: true)
      gss_api = OpenSSL::ASN1.decode(response.auth_value)
      security_blob = ::RubySMB::Gss.asn1dig(gss_api, 0, 2, 0)&.value
      ap_rep = Rex::Proto::Kerberos::Model::ApRep.decode(security_blob)
      ap_rep_enc_part = ap_rep.decrypt_enc_part(@session_key.value)
    rescue ::Rex::Proto::Kerberos::Model::Error::KerberosDecodingError,
             ::Rex::Proto::Kerberos::Model::Error::KerberosError,
             OpenSSL::ASN1::ASN1Error => e
      raise RubySMB::Dcerpc::Error::BindError, e.message # raise the more context-specific BindError
    end
    server_sequence_number = ap_rep_enc_part.sequence_number
    # Now complete the handshake - see [MS-KILE] 3.4.5.1 - https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-kile/190ab8de-dc42-49cf-bf1b-ea5705b7a087
    response_ap_rep = build_ap_rep(@session_key, server_sequence_number)

    wrapped_ap_rep = OpenSSL::ASN1::ASN1Data.new([
      OpenSSL::ASN1::Sequence.new([
        OpenSSL::ASN1::ASN1Data.new([
          OpenSSL::ASN1::OctetString(response_ap_rep.encode)
        ], 2, :CONTEXT_SPECIFIC)
      ])
    ], 1, :CONTEXT_SPECIFIC).to_der

    alter_ctx = RubySMB::Dcerpc::AlterContext.new(options)
    alter_ctx.pdu_header.call_id = @call_id

    add_auth_verifier(alter_ctx, wrapped_ap_rep)

    send_packet(alter_ctx)

    begin
      dcerpc_response = recv_struct(RubySMB::Dcerpc::AlterContextResp)
    rescue RubySMB::Dcerpc::Error::InvalidPacket
      raise RubySMB::Dcerpc::Error::BindError, e.message # raise the more context-specific BindError
    end

    self.krb_encryptor = @kerberos_authenticator.get_message_encryptor(ap_rep_enc_part.subkey,
                                                                                @client_sequence_number,
                                                                                server_sequence_number)
    # Set the session key value on the parent class - needed for decrypting attribute values in e.g. DRSR
    @session_key = ap_rep_enc_part.subkey.value
  end

  def get_auth_padding_length(plaintext_len)
    (16 - (self.krb_encryptor.calculate_encrypted_length(plaintext_len) % 16)) % 16
  end

  attr_accessor :krb_encryptor
end
